package com.mezz.bones.framework.websocket.handler;

import jakarta.websocket.Session;
import org.springframework.stereotype.Component;
import java.util.concurrent.*;

@Component
public class WebSocketHeartbeatHandler {
    private final ConcurrentHashMap<String, ScheduledFuture<?>> heartbeatTasks = new ConcurrentHashMap<>();
    private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);

    public void startHeartbeatTask(Session session, String heartBeatData) {
        String sessionId = session.getId(); // 使用Session的唯一ID
        
        ScheduledFuture<?> future = scheduler.scheduleAtFixedRate(() -> {
            try {
                if (session.isOpen()) {
                    session.getBasicRemote().sendText(heartBeatData);
                }
            } catch (Exception e) {
                stopHeartbeatTask(session);
            }
        }, 0, 30, TimeUnit.SECONDS);
        
        heartbeatTasks.put(sessionId, future);
    }
    
    public void stopHeartbeatTask(Session session) {
        ScheduledFuture<?> future = heartbeatTasks.remove(session.getId());
        if (future != null) {
            future.cancel(true);
        }
    }
}
